package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.feature.{HashingTF, IDF, IDFModel, Tokenizer}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo9TextClass {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[16]")
      .appName("point")
      .getOrCreate()
    import spark.implicits._
    import org.apache.spark.sql.functions._

    /**
     * 1、读取原始数据
     */
    val textDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", "\t")
      .schema("label DOUBLE,text STRING")
      .load("data/text.txt")

    textDF.show()

    /**
     * 2、分词
     */
    val ikDF: DataFrame = textDF.map {
      case Row(label: Double, text: String) =>
        //使用ik分词器分词
        val list: List[String] = Demo8IKUtils.fit(text)
        (label, list.mkString(" "))
    }.toDF("label", "text")

    /**
     * 3、清洗数据
     */
    val filterDF: DataFrame = ikDF
      .where($"text".isNotNull and $"text" =!= "")
      .where(size(split($"text", " ")) > 2)

    /**
     * 使用tf-idf将文本转换成向量
     */

    /**
     * 4、使用英文分词器分词
     */
    val tokenizer: Tokenizer = new Tokenizer()
      .setInputCol("text") //输入列
      .setOutputCol("words") //输出列
    val tokenizerDF: DataFrame = tokenizer.transform(filterDF)

    /**
     * 5、增加词频，每一个从再一段文本中出现的频率
     */
    val hashingTF: HashingTF = new HashingTF()
      .setInputCol("words") //输入列
      .setOutputCol("rawFeatures") //输出列
    //.setNumFeatures(20)//需要设置为总的词的数量

    val hashingTFDF: DataFrame = hashingTF.transform(tokenizerDF)

    /**
     * 6、增加IDF(逆文本频率)，用评估一个词普遍重要程度值
     */

    val idf: IDF = new IDF()
      .setInputCol("rawFeatures") //输入列
      .setOutputCol("features") //输出列
    val idfModel: IDFModel = idf.fit(hashingTFDF)

    val idfDF: DataFrame = idfModel.transform(hashingTFDF)


    /**
     * 7、将原始数据拆分成训练集和测试集
     */
    val Array(train, test) = idfDF.randomSplit(Array(0.8, 0.2))

    /**
     * 8、选择算法
     *
     */

    //贝叶斯分类
    val bayes = new NaiveBayes()

    /**
     * 9、将训练集带入算法训练模型
     *
     */

    val model: NaiveBayesModel = bayes.fit(train)

    /**
     * 10、使用测试集测试模型的准确率
     */
    val testDF: DataFrame = model.transform(test)


    /**
     * 计算准确率
     */
    val p: Double = testDF.where($"label" === $"prediction").count().toDouble / testDF.count()

    println(p)

  }

}
